from typing import TypedDict,Annotated

from langgraph.graph import add_messages,StateGraph,START,END

from langchain_openai import ChatOpenAI


class State(TypedDict):
    messages:Annotated[list,add_messages]

builder = StateGraph(State)

api_key = "sk-6S0PtpNia71gjcfwSsDPsJ9mGqsVPr2XRQzAx1dHbJS7RW4t"
api_base="https://chatapi.littlewheat.com/v1"

llm = ChatOpenAI(model="gpt-4o",api_key=api_key,base_url=api_base,temperature=0)

def chatbot(state:State):
    print(state)
    return  {"messages":[llm.invoke(state["messages"])]}

builder.add_node("chatbot",chatbot)
# builder.set_entry_point("chatbot")
# builder.set_finish_point("chatbot")
builder.add_edge(START,"chatbot")
builder.add_edge("chatbot",END)

graph = builder.compile()

# response = graph.invoke({"messages":"请介绍下自己"})
# print(response)

def stream_graph_updates(user_input:str):
    #for event in graph.stream({"messages":[{"user":user_input}]}):
    for event in graph.stream({"messages": [("user", user_input)]}):
        for value in event.values():
            print({"模型回复":value["messages"][-1].content})

while True:
    try:
        user_input = input("用户提问:")
        if user_input.lower() in ["退出"]:
            print("下次再见")
            break
        stream_graph_updates(user_input)
    except:
        break